import networkx as nx
import matplotlib.pyplot as plt

# 创建一个有向图
G = nx.DiGraph()

# 添加节点（代表不同时间点的分子）
G.add_node('A_t1', label='A', time='t1')
G.add_node('A_t2', label='A', time='t2')
G.add_node('B_t2', label='B', time='t2')
G.add_node('C_t3', label='C', time='t3')
G.add_node('D_t3', label='D', time='t3')

# 添加SAME_AS边（相同分子在不同时间点上的存在）
G.add_edge('A_t1', 'A_t2', relation='SAME_AS')

# 添加POTENTIAL_TRANSFORMATION边（潜在的化学转化）
G.add_edge('A_t2', 'B_t2', relation='POTENTIAL_TRANSFORMATION')
G.add_edge('B_t2', 'C_t3', relation='POTENTIAL_TRANSFORMATION')
G.add_edge('A_t2', 'D_t3', relation='POTENTIAL_TRANSFORMATION')

# 设置节点的颜色和大小
node_colors = []
node_sizes = []
for node in G.nodes(data=True):
    if node[1]['time'] == 't1':
        node_colors.append('lightblue')
        node_sizes.append(600)
    elif node[1]['time'] == 't2':
        node_colors.append('lightgreen')
        node_sizes.append(600)
    else:
        node_colors.append('lightcoral')
        node_sizes.append(600)

# 设置边的颜色和样式
edge_colors = []
edge_styles = []
for u, v, d in G.edges(data=True):
    if d['relation'] == 'SAME_AS':
        edge_colors.append('black')
        edge_styles.append('solid')
    elif d['relation'] == 'POTENTIAL_TRANSFORMATION':
        edge_colors.append('gray')
        edge_styles.append('dashed')

# 绘制图结构
pos = nx.spring_layout(G)  # 生成节点的位置
nx.draw(G, pos, with_labels=True, node_color=node_colors, node_size=node_sizes, edge_color=edge_colors, style=edge_styles, arrows=True)

# 显示图形
plt.show()
